import os
import pandas as pd
import numpy as np
from flask import Flask, render_template, request, redirect, url_for, flash, jsonify
from werkzeug.utils import secure_filename
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import seaborn as sns
from utils.data_handler import DataHandler
from utils.ml_models import MLModels
import io
import base64

app = Flask(__name__)
app.secret_key = 'your-secret-key-change-in-production'
app.config['UPLOAD_FOLDER'] = 'uploads'
app.config['MAX_CONTENT_LENGTH'] = 50 * 1024 * 1024  # 50MB

# 确保上传目录存在
os.makedirs(app.config['UPLOAD_FOLDER'], exist_ok=True)

# 全局变量存储数据（生产环境应使用数据库）
current_data = None
current_results = None

@app.route('/')
def index():
    return render_template('index.html')

@app.route('/upload', methods=['GET', 'POST'])
def upload():
    global current_data
    
    if request.method == 'POST':
        if 'file' not in request.files:
            flash('请选择文件')
            return redirect(request.url)
        
        file = request.files['file']
        if file.filename == '':
            flash('请选择文件')
            return redirect(request.url)
        
        if file and file.filename.endswith('.csv'):
            filename = secure_filename(file.filename)
            filepath = os.path.join(app.config['UPLOAD_FOLDER'], filename)
            file.save(filepath)
            
            try:
                # 读取数据
                current_data = pd.read_csv(filepath)
                data_handler = DataHandler(current_data)
                
                # 基础信息
                data_info = data_handler.get_basic_info()
                
                return render_template('analyze.html', 
                                     data_info=data_info,
                                     columns=current_data.columns.tolist(),
                                     data_sample=current_data.head().to_html(classes='table table-striped'))
            except Exception as e:
                flash(f'文件读取错误: {str(e)}')
                return redirect(request.url)
        else:
            flash('仅支持 CSV 格式文件')
            return redirect(request.url)
    
    return render_template('upload.html')

@app.route('/train', methods=['POST'])
def train():
    global current_data, current_results
    
    if current_data is None:
        flash('请先上传数据')
        return redirect(url_for('upload'))
    
    try:
        target_column = request.form.get('target_column')
        task_type = request.form.get('task_type')  # 'classification' or 'regression'
        
        if not target_column or target_column not in current_data.columns:
            flash('请选择有效的目标列')
            return redirect(url_for('upload'))
        
        # 训练模型
        ml_models = MLModels(current_data, target_column, task_type)
        results = ml_models.train_models()
        
        current_results = results
        
        return render_template('results.html', results=results)
        
    except Exception as e:
        flash(f'训练错误: {str(e)}')
        return redirect(url_for('upload'))

@app.route('/predict', methods=['POST'])
def predict():
    global current_results
    
    if current_results is None:
        return jsonify({'error': '请先训练模型'})
    
    try:
        # 获取输入特征
        input_data = {}
        for key, value in request.form.items():
            if key != 'model_name':
                input_data[key] = float(value)
        
        model_name = request.form.get('model_name', 'best_model')
        
        # 进行预测
        prediction = current_results['models'][model_name]['model'].predict([list(input_data.values())])
        
        return jsonify({
            'prediction': prediction[0],
            'model_used': model_name
        })
        
    except Exception as e:
        return jsonify({'error': str(e)})

if __name__ == '__main__':
    app.run(debug=True, host='0.0.0.0', port=5000)